"""
数据加载器
支持从Excel、JSON等文件加载测试数据
"""
import os
import json
import pandas as pd
from typing import Dict, List, Any, Union
from config.test_config import TestConfig
from utils.logger import Logger


class DataLoader:
    """数据加载器类"""
    
    def __init__(self):
        self.logger = Logger().get_logger()
        self.test_data_dir = TestConfig.TEST_DATA_DIR
    
    def load_excel_data(self, filename: str, sheet_name: str = None) -> List[Dict[str, Any]]:
        """
        从Excel文件加载数据
        
        Args:
            filename (str): Excel文件名
            sheet_name (str): 工作表名称，默认为第一个工作表
            
        Returns:
            List[Dict[str, Any]]: 数据列表
        """
        try:
            filepath = os.path.join(self.test_data_dir, filename)
            
            if not os.path.exists(filepath):
                raise FileNotFoundError(f"文件不存在: {filepath}")
            
            # 读取Excel文件
            if sheet_name:
                df = pd.read_excel(filepath, sheet_name=sheet_name)
            else:
                df = pd.read_excel(filepath)
            
            # 转换为字典列表
            data = df.to_dict('records')
            
            self.logger.info(f"成功加载Excel数据: {filename}, 共{len(data)}条记录")
            return data
            
        except Exception as e:
            self.logger.error(f"加载Excel数据失败: {str(e)}")
            raise
    
    def load_json_data(self, filename: str) -> Union[Dict[str, Any], List[Dict[str, Any]]]:
        """
        从JSON文件加载数据
        
        Args:
            filename (str): JSON文件名
            
        Returns:
            Union[Dict[str, Any], List[Dict[str, Any]]]: JSON数据
        """
        try:
            filepath = os.path.join(self.test_data_dir, filename)
            
            if not os.path.exists(filepath):
                raise FileNotFoundError(f"文件不存在: {filepath}")
            
            with open(filepath, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            self.logger.info(f"成功加载JSON数据: {filename}")
            return data
            
        except Exception as e:
            self.logger.error(f"加载JSON数据失败: {str(e)}")
            raise
    
    def load_csv_data(self, filename: str, encoding: str = 'utf-8') -> List[Dict[str, Any]]:
        """
        从CSV文件加载数据
        
        Args:
            filename (str): CSV文件名
            encoding (str): 文件编码
            
        Returns:
            List[Dict[str, Any]]: 数据列表
        """
        try:
            filepath = os.path.join(self.test_data_dir, filename)
            
            if not os.path.exists(filepath):
                raise FileNotFoundError(f"文件不存在: {filepath}")
            
            # 读取CSV文件
            df = pd.read_csv(filepath, encoding=encoding)
            
            # 转换为字典列表
            data = df.to_dict('records')
            
            self.logger.info(f"成功加载CSV数据: {filename}, 共{len(data)}条记录")
            return data
            
        except Exception as e:
            self.logger.error(f"加载CSV数据失败: {str(e)}")
            raise
    
    def get_test_data_by_key(self, filename: str, key: str, value: Any) -> Dict[str, Any]:
        """
        根据键值获取测试数据
        
        Args:
            filename (str): 数据文件名
            key (str): 查找的键
            value (Any): 查找的值
            
        Returns:
            Dict[str, Any]: 匹配的数据记录
        """
        try:
            # 根据文件扩展名选择加载方法
            if filename.endswith('.xlsx') or filename.endswith('.xls'):
                data = self.load_excel_data(filename)
            elif filename.endswith('.json'):
                data = self.load_json_data(filename)
                if isinstance(data, dict):
                    data = [data]
            elif filename.endswith('.csv'):
                data = self.load_csv_data(filename)
            else:
                raise ValueError(f"不支持的文件格式: {filename}")
            
            # 查找匹配的记录
            for record in data:
                if record.get(key) == value:
                    self.logger.info(f"找到匹配数据: {key}={value}")
                    return record
            
            raise ValueError(f"未找到匹配的数据: {key}={value}")
            
        except Exception as e:
            self.logger.error(f"获取测试数据失败: {str(e)}")
            raise
    
    def save_data_to_excel(self, data: List[Dict[str, Any]], filename: str, sheet_name: str = 'Sheet1'):
        """
        保存数据到Excel文件
        
        Args:
            data (List[Dict[str, Any]]): 要保存的数据
            filename (str): 文件名
            sheet_name (str): 工作表名称
        """
        try:
            filepath = os.path.join(self.test_data_dir, filename)
            
            # 转换为DataFrame
            df = pd.DataFrame(data)
            
            # 保存到Excel
            with pd.ExcelWriter(filepath, engine='openpyxl') as writer:
                df.to_excel(writer, sheet_name=sheet_name, index=False)
            
            self.logger.info(f"数据已保存到Excel: {filename}")
            
        except Exception as e:
            self.logger.error(f"保存Excel数据失败: {str(e)}")
            raise
    
    def save_data_to_json(self, data: Union[Dict[str, Any], List[Dict[str, Any]]], filename: str):
        """
        保存数据到JSON文件
        
        Args:
            data (Union[Dict[str, Any], List[Dict[str, Any]]]): 要保存的数据
            filename (str): 文件名
        """
        try:
            filepath = os.path.join(self.test_data_dir, filename)
            
            with open(filepath, 'w', encoding='utf-8') as f:
                json.dump(data, f, ensure_ascii=False, indent=2)
            
            self.logger.info(f"数据已保存到JSON: {filename}")
            
        except Exception as e:
            self.logger.error(f"保存JSON数据失败: {str(e)}")
            raise
    
    def get_all_test_files(self, file_type: str = None) -> List[str]:
        """
        获取所有测试数据文件
        
        Args:
            file_type (str): 文件类型过滤 (xlsx, json, csv)
            
        Returns:
            List[str]: 文件名列表
        """
        try:
            files = []
            
            if not os.path.exists(self.test_data_dir):
                return files
            
            for filename in os.listdir(self.test_data_dir):
                if file_type:
                    if filename.endswith(f'.{file_type}'):
                        files.append(filename)
                else:
                    if filename.endswith(('.xlsx', '.xls', '.json', '.csv')):
                        files.append(filename)
            
            return sorted(files)
            
        except Exception as e:
            self.logger.error(f"获取测试文件列表失败: {str(e)}")
            return [] 